MLflow Lab

%run "../Includes/Classroom-Setup"

from sklearn.model_selection import train_test_split

from sklearn.datasets import load_boston

from sklearn.preprocessing import StandardScaler

boston_housing = load_boston()

# split 80/20 train-test

X_train, X_test, y_train, y_test = train_test_split(boston_housing.data,

                                                    boston_housing.target,

                                                    test_size=0.2,

                                                    random_state=1)

# Scale features

scaler = StandardScaler()

X_train = scaler.fit_transform(X_train)

X_test = scaler.transform(X_test)

# Train-Validation Split

X_train_split, X_val, y_train_split, y_val = train_test_split(X_train,

                                                              y_train,

                                                              test_size=0.25,

                                                              random_state=1)

import tensorflow as tf

from tensorflow.keras.models import Sequential

from tensorflow.keras.layers import Dense

tf.random.set_seed(42)

def build_model():

  return Sequential([Dense(50, input_dim=13, activation="relu"),

                     Dense(20, activation="relu"),

                     Dense(1, activation="linear")])

# TODO

import mlflow

from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, LambdaCallback

filepath = f"{working_dir}/keras_mlflow.ckpt"

checkpointer = ModelCheckpoint(filepath=filepath, verbose=1, save_best_only=True)

earlyStopping = EarlyStopping(monitor="val_loss", min_delta=0.0001, patience=2, mode="auto")

mlflowCallback = LambdaCallback(on_epoch_end=lambda epoch, logs: mlflow.log_metric('loss', logs['loss']))

import matplotlib.pyplot as plt

def viewModelLoss(history):

  plt.clf()

  plt.semilogy(history.history["loss"], label="train_loss")

  plt.title("model loss")

  plt.ylabel("loss")

  plt.xlabel("epoch")

  plt.semilogy(history.history["val_loss"], label="val_loss")

  plt.legend()

  return plt

# TODO

from mlflow.keras import log_model

optimizer = "adam"

loss = "mse" 

epochs = 30

batch_size = 32

with mlflow.start_run() as run:

  model = build_model()

  model.compile(optimizer=optimizer, loss=loss, metrics=["mae", "mse"])

  mlflow.log_params({'loss': loss,

                    'optimizer': optimizer,

                    'epochs': epochs,

                    'batch_size': batch_size

                    })

  history = model.fit(X_train_split, 

                      y_train_split, 

                      validation_data=(X_val, y_val), 

                      epochs=epochs, 

                      batch_size=batch_size, 

                      callbacks=[checkpointer, earlyStopping, mlflowCallback], 

                      verbose=2)

  for i, layer in enumerate(model.layers):

    mlflow.log_param(f"hidden_layer_{i}_units", layer.output_shape)

  log_model(model, "keras_model")

  fig = viewModelLoss(history)

  fig.savefig("train-validation-loss.png")

  mlflow.log_artifact("train-validation-loss.png")

  plt.show()